#!/bin/bash

# Real-time scheduling priority setup script
# This script configures the system to allow regular users to use real-time scheduling

# Color definitions
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Ensure running as root
if [ "$EUID" -ne 0 ]; then
  echo -e "${RED}Error: Please run this script with root privileges (sudo $0)${NC}"
  exit 1
fi

# Set default values
TARGET_USER=$(logname 2>/dev/null || echo $SUDO_USER)
if [ -z "$TARGET_USER" ] || [ "$TARGET_USER" = "root" ]; then
  # Fallback methods to find the actual user
  TARGET_USER=$(who am i | awk '{print $1}')
  if [ -z "$TARGET_USER" ] || [ "$TARGET_USER" = "root" ]; then
    echo -e "${RED}Error: Unable to determine the current user${NC}"
    exit 1
  fi
fi

# Set default priority limits
SOFT_RTPRIO=20
HARD_RTPRIO=50

# Verify user exists
if ! id "$TARGET_USER" &>/dev/null; then
  echo -e "${RED}Error: User '$TARGET_USER' does not exist${NC}"
  exit 1
fi

echo -e "${BLUE}Configuring real-time scheduling limits for user '$TARGET_USER':${NC}"
echo -e "  Soft limit: ${SOFT_RTPRIO}"
echo -e "  Hard limit: ${HARD_RTPRIO}"

# Check limits.conf file
LIMITS_FILE="/etc/security/limits.conf"
if [ ! -f "$LIMITS_FILE" ]; then
  echo -e "${RED}Error: $LIMITS_FILE file does not exist${NC}"
  exit 1
fi

# Backup limits.conf
BACKUP_FILE="${LIMITS_FILE}.bak.$(date +%Y%m%d%H%M%S)"
cp "$LIMITS_FILE" "$BACKUP_FILE"
echo -e "${GREEN}Original configuration backed up to $BACKUP_FILE${NC}"

# Remove existing configuration if any
sed -i "/^${TARGET_USER}[[:space:]]\+\(soft\|hard\)[[:space:]]\+rtprio[[:space:]]\+/d" "$LIMITS_FILE"

# Add new configuration
echo -e "\n# Allow user $TARGET_USER to use real-time scheduling (added by setup_rtprio.sh on $(date))" >> "$LIMITS_FILE"
echo "$TARGET_USER soft rtprio $SOFT_RTPRIO" >> "$LIMITS_FILE"
echo "$TARGET_USER hard rtprio $HARD_RTPRIO" >> "$LIMITS_FILE"

echo -e "${GREEN}Successfully updated $LIMITS_FILE${NC}"

# Check PAM configuration and attempt to fix if needed
PAM_SESSION="/etc/pam.d/common-session"

# First check common locations
PAM_FOUND=false
for PAM_FILE in "/etc/pam.d/common-session" "/etc/pam.d/system-session" "/etc/pam.d/system-auth" "/etc/pam.d/login" "/etc/pam.d/su" "/etc/pam.d/sshd"; do
  if [ -f "$PAM_FILE" ]; then
    if grep -q "pam_limits.so" "$PAM_FILE"; then
      echo -e "${GREEN}Found pam_limits.so in $PAM_FILE${NC}"
      PAM_FOUND=true
      break
    else
      PAM_SESSION="$PAM_FILE"
    fi
  fi
done

# If not found in any file, try to add it to the most appropriate one
if [ "$PAM_FOUND" = false ]; then
  if [ -f "$PAM_SESSION" ]; then
    echo -e "${YELLOW}Adding pam_limits.so configuration to $PAM_SESSION${NC}"
    echo "session required pam_limits.so" >> "$PAM_SESSION"
    echo -e "${GREEN}Successfully updated PAM configuration${NC}"
  else
    echo -e "${YELLOW}Warning: Could not find a suitable PAM configuration file${NC}"
    echo -e "${YELLOW}Please manually ensure pam_limits.so is loaded in your PAM configuration${NC}"
  fi
fi

# Create test script - avoiding here-document to prevent syntax errors
TEST_DIR="/home/$TARGET_USER"
TEST_SCRIPT="$TEST_DIR/test_rtprio.py"

if [ -d "$TEST_DIR" ]; then
  echo '#!/usr/bin/env python3' > "$TEST_SCRIPT"
  echo 'import os' >> "$TEST_SCRIPT"
  echo 'import sys' >> "$TEST_SCRIPT"
  echo 'import resource' >> "$TEST_SCRIPT"
  echo '' >> "$TEST_SCRIPT"
  echo 'def check_rtprio():' >> "$TEST_SCRIPT"
  echo '    """Check rtprio limits and actual setting capability"""' >> "$TEST_SCRIPT"
  echo '    print("====== Real-time Scheduling Priority Test ======")' >> "$TEST_SCRIPT"
  echo '    ' >> "$TEST_SCRIPT"
  echo '    # Get current user' >> "$TEST_SCRIPT"
  echo '    try:' >> "$TEST_SCRIPT"
  echo '        username = os.getlogin()' >> "$TEST_SCRIPT"
  echo '    except:' >> "$TEST_SCRIPT"
  echo '        import pwd' >> "$TEST_SCRIPT"
  echo '        username = pwd.getpwuid(os.getuid())[0]' >> "$TEST_SCRIPT"
  echo '    print(f"Current user: {username}")' >> "$TEST_SCRIPT"
  echo '    ' >> "$TEST_SCRIPT"
  echo '    # Check rtprio limits' >> "$TEST_SCRIPT"
  echo '    try:' >> "$TEST_SCRIPT"
  echo '        soft, hard = resource.getrlimit(resource.RLIMIT_RTPRIO)' >> "$TEST_SCRIPT"
  echo '        print(f"RTPRIO limits: soft={soft}, hard={hard}")' >> "$TEST_SCRIPT"
  echo '    except AttributeError:' >> "$TEST_SCRIPT"
  echo '        print("Unable to get RTPRIO limits (Python may lack support)")' >> "$TEST_SCRIPT"
  echo '    ' >> "$TEST_SCRIPT"
  echo '    # Try setting real-time scheduling' >> "$TEST_SCRIPT"
  echo '    print("\\nTrying different priority levels:")' >> "$TEST_SCRIPT"
  echo '    test_priorities = [1, 5, 10, 20, 30, 50, 80]' >> "$TEST_SCRIPT"
  echo '    ' >> "$TEST_SCRIPT"
  echo '    for priority in test_priorities:' >> "$TEST_SCRIPT"
  echo '        try:' >> "$TEST_SCRIPT"
  echo '            os.sched_setscheduler(0, os.SCHED_RR, os.sched_param(priority))' >> "$TEST_SCRIPT"
  echo '            policy = os.sched_getscheduler(0)' >> "$TEST_SCRIPT"
  echo '            policy_name = {' >> "$TEST_SCRIPT"
  echo '                0: "SCHED_OTHER (normal)",' >> "$TEST_SCRIPT"
  echo '                1: "SCHED_FIFO (real-time FIFO)",' >> "$TEST_SCRIPT"
  echo '                2: "SCHED_RR (real-time round-robin)"' >> "$TEST_SCRIPT"
  echo '            }.get(policy, f"Unknown ({policy})")' >> "$TEST_SCRIPT"
  echo '            ' >> "$TEST_SCRIPT"
  echo '            print(f"√ Priority {priority}: Successfully set to {policy_name}")' >> "$TEST_SCRIPT"
  echo '        except PermissionError:' >> "$TEST_SCRIPT"
  echo '            print(f"× Priority {priority}: Permission denied")' >> "$TEST_SCRIPT"
  echo '        except Exception as e:' >> "$TEST_SCRIPT"
  echo '            print(f"× Priority {priority}: Failed - {str(e)}")' >> "$TEST_SCRIPT"
  echo '    ' >> "$TEST_SCRIPT"
  echo '    print("\\nTest completed!")' >> "$TEST_SCRIPT"
  echo '' >> "$TEST_SCRIPT"
  echo 'if __name__ == "__main__":' >> "$TEST_SCRIPT"
  echo '    check_rtprio()' >> "$TEST_SCRIPT"

  chown "$TARGET_USER":"$TARGET_USER" "$TEST_SCRIPT"
  chmod +x "$TEST_SCRIPT"

  echo -e "${GREEN}Test script created: $TEST_SCRIPT${NC}"
else
  echo -e "${YELLOW}Warning: Cannot create test script (user home directory not found)${NC}"
fi

# Display usage instructions
echo
echo -e "${BLUE}==== Configuration Complete ====${NC}"
echo -e "${BLUE}Notes:${NC}"
echo -e "1. User ${GREEN}$TARGET_USER${NC} must log out and log back in for settings to take effect"
echo -e "2. You can verify the configuration with the test script:"
echo -e "   ${YELLOW}su - $TARGET_USER -c \"python3 $TEST_SCRIPT\"${NC}"
echo
echo -e "Typical Python code to use real-time scheduling:"
echo -e "${YELLOW}import os"
echo -e "os.sched_setscheduler(0, os.SCHED_RR, os.sched_param($SOFT_RTPRIO))${NC}"
echo
echo -e "${BLUE}Important:${NC}"
echo -e "- Use real-time priorities carefully to avoid system unresponsiveness"
echo -e "- Consider adding timeout protection in your programs"
echo -e "- To restore default configuration, edit ${GREEN}$LIMITS_FILE${NC} and remove the related lines"
echo

exit 0
